# GPL v2
# Todd Korody 2007
# 1 = mp3, 2 = ogg, 3 = flac
#TODO Make flac/ogg not have dup track tags (optional perhaps?)

import sys, base64, mimetypes
supported_File_Types = []

# Lets try to import all we can.  What we can import will be added to the supported list
# that is available from a static method.
try:
	import pyid3lib
	supported_File_Types.append('mp3')
	supported_File_Types.append('audio/mpeg')
except ImportError:
	print "You lack mp3 support"
try:
	import ogg.vorbis
	supported_File_Types.append('ogg')
	supported_File_Types.append('application/ogg')
	# legacy
	supported_File_Types.append('application/x-ogg')
except ImportError:
	print "You lack ogg support"
try:
	import flac.metadata as metadata
	supported_File_Types.append('flac')
	supported_File_Types.append('application/x-flac')
except ImportError:
	print "You lack flac support"

class FileNotSupported(Exception):
	pass
class InvalidInput(Exception):
	pass

class tags:
	def __init__(self, filename):
		self.filename = filename
		DEBUG = 0
		filetype = mimetypes.guess_type(filename)[0]
		print filetype
#		if supported_File_Types.count(file[-3:]) < 1 and  supported_File_Types.count(file[-4:]) < 1:
#			raise FileNotSupported, str(file) + ' is not of supported type'
		if filetype not in supported_File_Types:
			raise FileNotSupported, str(file) + ' is not of supported type'
		if filename[-4:] == '.mp3':
			self.type = 1
			try:
				self.tag = pyid3lib.tag( filename )
			except:
				raise FileNotSupported, str(filename) + ' could not be opened'

			try:
				self.album = [self.tag.album]
			except AttributeError:
				self.album = []
			try:
				self.artist = [self.tag.artist]
			except AttributeError:
				self.artist = []

			self.vendor = []

			try:
				self.title = [self.tag.title]
			except AttributeError:
				self.title = []
			try:
				self.tracktotal = [self.tag.track[1]]
			except AttributeError:
				self.tracktotal = []
			try:
				self.tracknumber = [self.tag.track[0]]
			except AttributeError:
				self.tracknumber = []

			try:
				self.genre = [self.tag[self.tag.index('TCON')]['text']]
			except ValueError:
				self.genre = []

			try:
				self.date = [self.tag.year]
			except AttributeError:
				self.date = []
			try:
				self.discnum = [self.tag[self.tag.index('TPOS')]['text']]
			except ValueError:
				self.discnum = []
			try:
				self.comment = [self.tag[self.tag.index('COMM')]['text']]
			except ValueError:
				self.comment = []
			if DEBUG == 1:
				for i in self.tag:
					print i

		elif filename[-4:] == '.ogg':
			self.type = 2
			try:
				self.tag = ogg.vorbis.VorbisFile(filename)
			except ogg.vorbis.VorbisError:
				raise FileNotSupported, str(filename) + ' is not of supported type'				

			if DEBUG == 1:
				print self.tag.comment().as_dict()

			self.album = self.get_Ogg_Field_List('ALBUM')
			self.artist = self.get_Ogg_Field_List('ARTIST')
			self.vendor = self.get_Ogg_Field_List('VENDOR')
			self.title = self.get_Ogg_Field_List('TITLE')
			self.tracktotal = self.get_Ogg_Field_List('TRACKTOTAL')
			self.tracknumber = self.get_Ogg_Field_List('TRACKNUMBER')
			self.genre = self.get_Ogg_Field_List('GENRE')
			self.discnum = self.get_Ogg_Field_List('DISCNUMBER')
			self.date = self.get_Ogg_Field_List('DATE')
			self.comment = self.get_Ogg_Field_List('DESCRIPTION')
				
		elif filename[-5:] == '.flac':
			self.type = 3
			# MUCH of the following is taken from the python-flac example code
			# create a chain
			chain = metadata.Chain()
			chain.read(filename)

			# get iterator, init
			it = metadata.Iterator()
			it.init(chain)

			while 1:
    				if it.get_block_type() == metadata.VORBIS_COMMENT:
					block = it.get_block()
					vc = metadata.VorbisComment(block)
					break
				if not it.next():
					break

			if not vc:
				raise FileNotSupported, str(filename) + ' does not seem to have a comment block'
			if DEBUG == 1:
				for c in vc.comments:
					print c

			self.chain = chain
			self.vc = vc

			self.album = self.get_Flac_Field_List('ALBUM')
			self.artist = self.get_Flac_Field_List('ARTIST')
			self.vendor = self.get_Flac_Field_List('VENDOR')
			self.title = self.get_Flac_Field_List('TITLE')
			self.tracktotal = self.get_Flac_Field_List('TRACKTOTAL')
			self.tracknumber = self.get_Flac_Field_List('TRACKNUMBER')
			self.genre = self.get_Flac_Field_List('GENRE')
			self.discnum = self.get_Flac_Field_List('DISCNUMBER')
			self.date = self.get_Flac_Field_List('DATE')
			self.comment = self.get_Flac_Field_List('DESCRIPTION')



	def get_Supported_Formats():
		return supported_File_Types
	get_Supported_Formats = staticmethod(get_Supported_Formats)

	def get_Album(self):
			return self.album
	def get_Artist(self):
			return self.artist
	def get_Vendor(self):
			return self.vendor # ERROR
	def get_Title(self):
			return self.title
	def get_TrackTotal(self):
			return self.tracktotal
	def get_TrackNumber(self):
			return self.tracknumber
	def get_Genre(self):
			return self.genre
	def get_DiscNumber(self):
			return self.discnum
	def get_Date(self):
			return self.date
	def get_Comment(self):
			return self.comment

	def set_Album(self, newalbum):
		if self.type == 1:
			self.tag.album = newalbum
			self.tag.update()
			self.album = [newalbum]
		elif self.type == 2:
			if self.write_Ogg_Tag('ALBUM',newalbum) == True:
				self.album.append(unicode(newalbum))
		elif self.type == 3:
			if self.write_Flac_Tag('ALBUM',newalbum) == True:
				self.album.append(unicode(newalbum))

	def set_Artist(self, newartist):
		if self.type == 1:
			self.tag.artist = newartist
			self.tag.update()
			self.artist = [newartist]
		elif self.type == 2:
			if self.write_Ogg_Tag('ARTIST',newartist) == True:
				self.artist.append(unicode(newartist))
			else:
				print "DUPS NO CHANGES"
		elif self.type == 3:
			if self.write_Flac_Tag('ARTIST',newartist) == True:
				self.artist.append(unicode(newartist))

	def set_Title(self, newtitle):
		if self.type == 1:
			self.tag.title = newtitle
			self.tag.update()
			self.title = [newtitle]
		elif self.type == 2:
			if self.write_Ogg_Tag('TITLE',newtitle) == True:
				self.title.append(unicode(newtitle))
		elif self.type == 3:
			if self.write_Flac_Tag('TITLE',newtitle) == True:
				self.title.append(unicode(newtitle))

	def set_TrackTotal(self, newtracktotal):
		if self.type == 1:
			if self.tracknumber == []:
				a = 0
			else:
				a = self.tracknumber[0]
			self.tag.track = (int(a), int(newtracktotal))
			self.tag.update()
			self.tracktotal = [newtracktotal]
		elif self.type == 2:
			if self.write_Ogg_Tag('TRACKTOTAL',newtracktotal) == True:
				self.tracktotal.append(unicode(newtracktotal))
		elif self.type == 3:
			if self.write_Flac_Tag('TRACKTOTAL',newtracktotal) == True:
				self.tracktotal.append(unicode(newtracktotal))

	def set_TrackNumber(self, newtracknumber):
		if self.type == 1:
			if self.tracktotal == []:
				a = 0
			else:
				a = self.tracktotal[0]
			self.tag.track = (int(newtracknumber), int(a))
			self.tag.update()
			self.tracknumber = [newtracknumber]
		elif self.type == 2:
			if self.write_Ogg_Tag('TRACKNUMBER',newtracknumber) == True:
				self.tracknumber.append(unicode(newtracknumber))
		elif self.type == 3:
			if self.write_Flac_Tag('TRACKNUMBER',newtracknumber) == True:
				self.tracknumber.append(unicode(newtracknumber))

	def set_Genre(self, newgenre):
		if self.type == 1:
			self.set_Mp3_Fields('TCON', newgenre)
			self.genre = [newgenre]
		elif self.type == 2:
			if self.write_Ogg_Tag('GENRE',newgenre) == True:
				self.genre.append(unicode(newgenre))
		elif self.type == 3:
			if self.write_Flac_Tag('GENRE',newgenre) == True:
				self.genre.append(unicode(newgenre))

	def set_DiscNumber(self, newdiscnum):
		if self.type == 1:
			self.set_Mp3_Fields('TPOS', newdiscnum)
			self.discnum = [newdiscnum]
		elif self.type == 2:
			if self.write_Ogg_Tag('DISCNUMBER',newdiscnum) == True:
				self.discnum.append(unicode(newdiscnum))
		elif self.type == 3:
			if self.write_Flac_Tag('DISCNUMBER',newdiscnum) == True:
				self.discnum.append(unicode(newdiscnum))

	def set_Date(self, newdate):
		if len(str(newdate)) != 4:		
			raise InvalidInput, ' The Year must be 4 digits'	
		if self.type == 1:
			self.tag.year = newdate
			self.tag.update()
			self.date = [newdate]
		elif self.type == 2:
			if self.write_Ogg_Tag('DATE',newdate) == True:
				self.date.append(unicode(newdate))
		elif self.type == 3:
			if self.write_Flac_Tag('DATE',newdate) == True:
				self.date.append(unicode(newdate))

	def set_Comment(self, newcomment):
		if self.type == 1:
			self.set_Mp3_Fields('COMM', newcomment)
			self.comment = [newcomment]
		elif self.type == 2:
			if self.write_Ogg_Tag('DESCRIPTION',newcomment) == True:
				self.comment.append(unicode(newcomment))
		elif self.type == 3:
			if self.write_Flac_Tag('DESCRIPTION',newcomment) == True:
				self.comment.append(unicode(newcomment))

	def get_AlbumArt(self, picturetype=None):
		if self.type == 1:
			# a list containing all matches
			returnval = []
			# default to "front cover" it just seems obvious
			if picturetype == None:
				picturetype = 3
			a = self.get_Mp3_Fields('APIC')
			for x in a:
				if x['picturetype'] == picturetype:
					return (x['mimetype'],x['data'])
			# No Album Art Matching The Restrictions Found
			return (None,None)

		elif self.type == 2:
			blob = self.get_Ogg_Field('COVERART')
			mimetype = self.get_Ogg_Field('COVERARTMIME')
			# it would be bad if the 'or' actually mattered here
			if mimetype == None or blob == None:			
				return (None,None)
			text = base64.decodestring(blob)
			return (mimetype,text)
		elif self.type == 3:
			blob = self.get_Flac_Field('COVERART')
			mimetype = self.get_Flac_Field('COVERARTMIME')
			# it would be bad if the 'or' actually mattered here
			if mimetype == None or blob == None:			
				return (None,None)
			text = base64.decodestring(blob)
			return (mimetype,text)

	def set_AlbumArt(self, filename, description='Album Art', picturetype=3):
		try:
			filetype = mimetypes.guess_type(filename)[0]
		except TypeError:
			raise FileNotSupported, str(filename) + ' could not pass mime inspection'

		if filetype != 'image/jpeg' and filetype != 'image/png':
			raise FileNotSupported, str(filename) + ' is not of supported type for album art'	
		try:
			filename = open(filename, 'rb')
			filedata = filename.read()
			filename.close()
		except:
			raise FileNotSupported, str(filename) + ' could not be read'	

		if self.type == 1:
			frames = []
			# make a list of all picture frames matching the picturetype
			for i in self.tag:
				if i['frameid'] == 'APIC' and i['picturetype'] == int(picturetype):
					frames.append(i)
			# del all those frames
			for i in frames:
				self.tag.remove( i['frameid'] )

			# now we are sure there are no dups so insert away
			d = { 'frameid' : 'APIC', 'mimetype' : filetype,
			'description' : description,'picturetype' : int(picturetype),'data' : filedata }
			self.tag.append(d)
			self.tag.update()
		elif self.type == 2:
			b = base64.encodestring(filedata)
			self.write_Ogg_Tag('COVERARTMIME',filetype)
			self.write_Ogg_Tag('COVERART',b)
		elif self.type == 3:
			b = base64.encodestring(filedata)
			self.write_Flac_Tag('COVERARTMIME',filetype)
			self.write_Flac_Tag('COVERART',b)

	def get_Ogg_Field(self, field):
		values = self.tag.comment().as_dict()
		try:
			return values[field].pop()
		except KeyError:
			return []

	def get_Ogg_Field_List(self, field):
		values = self.tag.comment().as_dict()
		try:
			return values[field]
		except KeyError:
			return []

	def get_Flac_Field1(self, field):
		for comment in self.vc.comments:
			b = comment.split('=')
			if b[0] == field:
				actualdata = comment.replace(field + '=','')
				return actualdata
		# What we are looking for must not exist
		return []

	def get_Flac_Field_List(self, field):
		returnval = []
		for comment in self.vc.comments:
			b = comment.split('=')
			if b[0] == field:
				actualdata = str(comment).replace(field + '=','')
				returnval.append(actualdata)
		return returnval
#		if len(returnval) == 0:
#			return None
#		else:
#			return returnval

	def get_Flac_Field(self, field):
		try:
			return self.vc.comments[field]
		except KeyError:
			return []

	def write_Ogg_Tag(self, field, newdata):
		comment = self.tag.comment()
		try:
			if str(newdata).upper() in str(comment.as_dict()[field]).upper():
				# The Data Is Already In The Tag
				return False
		except KeyError:
			# The File Doesn't Yet Have This Field
			pass
		# Add The Field
		comment.add_tag(field, unicode(newdata))
		# Write Out The File
		comment.write_to(self.filename)
		return True

	def write_Flac_Tag(self, field, newdata):
		does_exist = field + '=' + str(newdata).strip()
			# This loop checks to see if the tag already exists
		for c in self.vc.comments:
			if c.upper() == does_exist.upper():
				return False
		self.vc.comments[field] = newdata
		self.chain.write(True,True)
		return True

	def clear_Tag_Data(self):
		if self.type == 1:
			frames = []
			# make a list of all frames
			for i in self.tag:
				frames.append(i)
			# del all those frames
			for i in frames:
				self.tag.remove( i['frameid'] )
			self.tag.update()
		elif self.type == 2:
			self.tag.comment().clear()
			self.tag.comment().write_to(self.filename)
		elif self.type == 3:
			self.vc.block.vorbiscomment_remove_entries_matching('ARTIST')
			self.vc.block.vorbiscomment_remove_entries_matching('ALBUM')
			self.vc.block.vorbiscomment_remove_entries_matching('TITLE')
			self.vc.block.vorbiscomment_remove_entries_matching('DATE')
			self.vc.block.vorbiscomment_remove_entries_matching('DISCNUMBER')
			self.vc.block.vorbiscomment_remove_entries_matching('DESCRIPTION')
			self.vc.block.vorbiscomment_remove_entries_matching('TRACKTOTAL')
			self.vc.block.vorbiscomment_remove_entries_matching('TRACKNUMBER')
			self.vc.block.vorbiscomment_remove_entries_matching('GENRE')
			self.vc.block.vorbiscomment_remove_entries_matching('COVERARTMIME')			
			self.vc.block.vorbiscomment_remove_entries_matching('COVERART')

			for c in self.vc.comments:
				print c
			self.chain.write(True,True)
			
		self.album = []
		self.artist = []
		self.vendor = []
		self.title = []
		self.tracktotal = []
		self.tracknumber = []
		self.genre = []
		self.discnum = []
		self.date = []
		self.comment = []

	def set_Mp3_Fields(self, field, newdata):
		try:
			d = self.tag[self.tag.index(field)]        # get the frame
			d['text'] = str(newdata)       # update the dictionary
			self.tag[self.tag.index(field)] = d  # write back the dictionary
		except ValueError:
			# the item doesn't exist so just go ahead and add it
			self.tag.append( { 'frameid' : field, 'text' : str(newdata) } )
		self.tag.update()

	def get_Mp3_Fields(self, field):
		fields = []
		for x in self.tag:
			if x['frameid'] == field:
				fields.append(x)
		return fields
